from PendulumDLL import *
import gymnasium as gym
import torch
from agents import Actor
from tqdm import tqdm


def eval_agent(actor_path: str, eval_env: gym.Env):
	actor = Actor(eval_env.observation_space.shape[0], eval_env.action_space.shape[0], [64, 32]).to().cuda()
	actor.load_state_dict(torch.load(actor_path))
	max_step = eval_env.max_step
	state, _ = eval_env.reset()
	reward_sum = 0
	for _ in tqdm(range(max_step)):
		state_TEN = torch.from_numpy(state).unsqueeze(0).to().cuda()
		action = actor(state_TEN, eval=True)[0].squeeze(0).cpu().detach().numpy()
		state, reward, done, truncated, info = eval_env.step(action)
		reward_sum += reward
		if done: break
		print(f"| theta: {info['theta']:.2f}")
	print(f"| Reward sum: {reward_sum:.2f}")
	eval_env.close()


if __name__ == '__main__':
	actor_path = './train_logs/8246__Reward=-0.000__actor.pth'
	eval_env = gym.make('PendulumDLL-v0')
	eval_agent(actor_path, eval_env)
